package com.jzy.chatglmsdk18753goldn.test;

import ai.z.openapi.service.model.ChatMessageRole;
import com.alibaba.fastjson.JSON;
import com.jzy.chatglmsdk18753goldn.model.EventType;
import com.jzy.chatglmsdk18753goldn.model.Model;
import com.jzy.chatglmsdk18753goldn.model.ChatCompletionCommonRequest;
import com.jzy.chatglmsdk18753goldn.model.ChatCompletionResponse;
import com.jzy.chatglmsdk18753goldn.model.GLM.GLMChatCompletionRequest;
import com.jzy.chatglmsdk18753goldn.session.Configuration;
import com.jzy.chatglmsdk18753goldn.session.OpenAiSession;
import com.jzy.chatglmsdk18753goldn.session.OpenAiSessionFactory;
import com.jzy.chatglmsdk18753goldn.session.defaults.DefaultOpenAiSessionFactory;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.logging.HttpLoggingInterceptor;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.Nullable;
import org.junit.Before;
import org.junit.Test;

import java.util.Arrays;
import java.util.concurrent.CountDownLatch;

@Slf4j
public class GLM45Test {

    private OpenAiSession openAiSession;

    @Before
    public void init() {
        // 1. 配置文件
        Configuration configuration = new Configuration();
        configuration.setApiHost("https://open.bigmodel.cn/");
        configuration.setApiSecretKey_GLM("bb70cc6355ef46cdb890e7b5359cf9c2.qD9j4XV15Dwa1kD8");
        configuration.setLevel(HttpLoggingInterceptor.Level.BODY);
        configuration.setApiHost("https://dashscope.aliyuncs.com/");
        configuration.setApiSecretKey_Qwen("sk-a53f260c382b48fda4519294cbe91284");
        // 2. 会话工厂
        OpenAiSessionFactory factory = new DefaultOpenAiSessionFactory(configuration);
        // 3. 开启会话
        this.openAiSession = factory.openSession();
    }

    @Test
    public void testGLM45() throws Exception {
        CountDownLatch countDownLatch = new CountDownLatch(1);

        GLMChatCompletionRequest request = new GLMChatCompletionRequest();
        request.setModel(Model.GLM_45_FLASH);
        request.setMessages(Arrays.asList(
                ChatCompletionCommonRequest.Prompt.builder()
                        .role(ChatMessageRole.USER.value())
                        .content("你好")
                        .build()
        ));
        request.setStream(true);
        request.setTemperature(0.7f);
        request.setThinking(GLMChatCompletionRequest.Think.disabled());
        request.setMaxTokens(1024);

        System.out.println(request.toString());

        openAiSession.completions(request, new EventSourceListener() {
            @Override
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                ChatCompletionResponse response = JSON.parseObject(data, ChatCompletionResponse.class);
                log.info("测试结果 onEvent：{}", response.getData());
                // type 消息类型，add 增量，finish 结束，error 错误，interrupted 中断
                if (EventType.finish.getCode().equals(type)) {
                    ChatCompletionResponse.Meta meta = JSON.parseObject(response.getMeta(), ChatCompletionResponse.Meta.class);
                    log.info("[输出结束] Tokens {}", JSON.toJSONString(meta));
                }
            }

            @Override
            public void onClosed(EventSource eventSource) {
                log.info("对话完成");
                countDownLatch.countDown();
            }

            @Override
            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                log.info("对话异常");
                countDownLatch.countDown();
            }
        });
        countDownLatch.await();
    }

    @Test
    public void testGLMZ1() throws Exception {
        CountDownLatch countDownLatch = new CountDownLatch(1);

        ChatCompletionCommonRequest request = ChatCompletionCommonRequest.builder()
                .model(Model.GLM_Z1_FLASH)
                .messages(Arrays.asList(
                        ChatCompletionCommonRequest.Prompt.builder()
                                .role(ChatMessageRole.USER.value())
                                .content("你好")
                                .build()
                ))
                .stream(true)
                .temperature(0.7f)
                .build();

        System.out.println(request.toString());

        openAiSession.completions(request, new EventSourceListener() {
            @Override
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                ChatCompletionResponse response = JSON.parseObject(data, ChatCompletionResponse.class);
                log.info("测试结果 onEvent：{}", response.getData());
                // type 消息类型，add 增量，finish 结束，error 错误，interrupted 中断
                if (EventType.finish.getCode().equals(type)) {
                    ChatCompletionResponse.Meta meta = JSON.parseObject(response.getMeta(), ChatCompletionResponse.Meta.class);
                    log.info("[输出结束] Tokens {}", JSON.toJSONString(meta));
                }
            }

            @Override
            public void onClosed(EventSource eventSource) {
                log.info("对话完成");
                countDownLatch.countDown();
            }

            @Override
            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                log.info("对话异常");
                countDownLatch.countDown();
            }
        });
        countDownLatch.await();
    }

}
